package com.oak.broke.config.mybatis.encript.interceptor;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLInListExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
import com.alibaba.druid.util.JdbcConstants;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.oak.broke.config.mybatis.encript.encrypt.ProcessEncrypt;
import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.lang.reflect.Proxy;
import java.sql.PreparedStatement;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;

@Intercepts({@Signature(
        type = ParameterHandler.class,
        method = "setParameters",
        args = {PreparedStatement.class}
)})
@Component
public class EncryptionInterceptor implements Interceptor {

    @Autowired
    private ProcessEncrypt processEncrypt;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        ParameterHandler parameterHandler = PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(parameterHandler);
        BoundSql boundSql = (BoundSql) metaObject.getValue("boundSql");
        String sql = boundSql.getSql();
        Map<Integer, String> indexToColumnMap = getIndexToColumnMap(sql);
        PreparedStatement ps = (PreparedStatement) invocation.getArgs()[0];
        PreparedStatement psProxy = (PreparedStatement) Proxy.newProxyInstance(PreparedStatement.class.getClassLoader()
                , new Class[]{PreparedStatement.class}
                , (proxy, method, args) -> {
                    if (method.getName().startsWith("set")) {
                        //PreparedStatement 下标从1开始
                        int index = ((int) args[0]) - 1;
                        Object value = args[1];
                        //替换参数为加密过后的参数
                        args[1] = processEncrypt.processEncrypt(indexToColumnMap.get(index), value);
                    }
                    return method.invoke(ps, args);
                });
        return invocation.getMethod().invoke(invocation.getTarget(), psProxy);
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
    }

    private static Map<Integer, String> getIndexToColumnMap(String sql) {
        final Map<Integer, String> indexToColumnMap = new LinkedHashMap<Integer, String>();
        List<SQLStatement> stmts = SQLUtils.parseStatements(sql, JdbcConstants.MYSQL);
        final int[] index = {0};
        for (SQLStatement stmt : stmts) {
            MySqlASTVisitorAdapter sqlastVisitorAdapter = new MySqlASTVisitorAdapter() {
                @Override
                public boolean visit(SQLBinaryOpExpr x) {
                    if (SQLBinaryOperator.Equality.equals(x.getOperator())) {
                        if (x.getRight() instanceof SQLVariantRefExpr) {
                            if ("?".equals(((SQLVariantRefExpr) x.getRight()).getName())) {
                                indexToColumnMap.put(index[0], getColName(x.getLeft()));
                                index[0]++;
                            }
                        }
                    }
                    // ...其他条件自行添加
                    return super.visit(x);
                }

                @Override
                public boolean visit(SQLUpdateSetItem x) {
                    if ("?".equals(x.getValue().toString())) {
                        indexToColumnMap.put(index[0], getColName(x.getColumn()));
                        index[0]++;
                    }
                    return super.visit(x);
                }

                @Override
                public boolean visit(SQLInListExpr x) {
                    for (SQLExpr expr : x.getTargetList()) {
                        if (expr instanceof SQLVariantRefExpr) {
                            if ("?".equals(((SQLVariantRefExpr) expr).getName())) {
                                indexToColumnMap.put(index[0], getColName(x.getExpr()));
                                index[0]++;
                            }

                        }
                    }
                    return super.visit(x);
                }

                @Override
                public boolean visit(MySqlInsertStatement x) {
                    List<SQLExpr> columns = x.getColumns();
                    List<SQLInsertStatement.ValuesClause> valuesClauses = x.getValuesList();
                    for (SQLInsertStatement.ValuesClause valuesClause : valuesClauses) {
                        List<SQLExpr> values = valuesClause.getValues();
                        for (int i = 0, size = values.size(); i < size; i++) {
                            if (values.get(i) instanceof SQLVariantRefExpr) {
                                if ("?".equals(((SQLVariantRefExpr) values.get(i)).getName())) {
                                    indexToColumnMap.put(index[0], getColName(columns.get(i)));
                                    index[0]++;
                                }
                            }
                        }
                    }
                    return super.visit(x);
                }
            };
            stmt.accept(sqlastVisitorAdapter);
        }
        return indexToColumnMap;
    }

    private static String getColName(SQLExpr exp) {
        String name = null;
        if (exp instanceof SQLPropertyExpr) {
            name = ((SQLPropertyExpr) exp).getName();
        } else if (exp instanceof SQLIdentifierExpr) {
            name = ((SQLIdentifierExpr) exp).getName();
        }
        return name;
    }
}
